from chapter_2_2_DMA import gen_dma_strategy_data
import matplotlib.pyplot as plt
import pandas as pd

dma_strategy_data = gen_dma_strategy_data('002419', '20220101', '20221231')

print(dma_strategy_data)

'''
strategy_data: 策略的数据
initial_cash: 初始资金
'''

def draw_dma_stock_cash(data):
    plt.figure(figsize=(10, 6))
    plt.plot(data['stock_value'], lw=2, label='stock_value')
    plt.plot(data['total'], lw=2, ls='--', label='total')
    plt.legend()
    plt.grid()
    plt.show()


def test(strategy_data, initial_cash):
    print('============test DMA===============')
    # 新建一个positions表，序号和strategy数据表保持一致
    positions = pd.DataFrame(index=strategy_data.index).fillna(0)
    print(positions)

    # 因为A股都是最低100股
    # 因此设置stock字段为交易信号的100倍
    positions['stock'] = strategy_data['signal'] * 100
    portfolio = pd.DataFrame(index=positions.index)
    portfolio['stock_value'] = positions.multiply(strategy_data['price'], axis=0)
    # 仓位变价就是下单的数量
    order = positions.diff()
    print('交易量')
    print(order.tail(20))
    portfolio['stock'] = positions['stock']
    # 初始资金减去下单金额的总和就是剩余的资金
    portfolio['cash'] = initial_cash - order.multiply(strategy_data['price'], axis=0).cumsum()
    portfolio['total'] = portfolio['cash'] + portfolio['stock_value']
    print(portfolio.tail(10))
    draw_dma_stock_cash(portfolio)


test(dma_strategy_data, 20000)


# 绘图
def draw_dma_chart(data):
    # 绘制
    plt.figure(figsize=(50, 12))
    # 绘制股价变化
    plt.plot(data['price'], lw=1, c='y', label="price")
    plt.plot(data['avg_5'], ls='--', c='r', lw=2, label="avg5")
    plt.plot(data['avg_10'], ls='-.', c='b', lw=2, label="avg10")

    # 买入卖出信号标记
    plt.scatter(
        data.loc[data.order == 1].index,
        data['price'][data.order == 1],
        marker='^', s=50, color='r', label='buy'
    )
    plt.scatter(
        data.loc[data.order == -1].index,
        data['price'][data.order == -1],
        marker='v', s=50, color='g', label='sell'
    )

    plt.legend()
    plt.grid()
    plt.show()

draw_dma_chart(dma_strategy_data)
